from dotenv import load_dotenv
import os
from openai import OpenAI

# 从.env文件中获取API密钥
def get_api_key():
    load_dotenv()  # 加载环境变量文件
    api_key = os.getenv("DEEPSEEK_API_KEY", "")  # 获取DEEPSEEK_API_KEY环境变量
    if not api_key:
        print("DEEPSEEK_API_KEY is not set in the.env file")
        return None  # 如果密钥未设置则返回None
    return api_key

# 根据平台获取API基础URL
def get_base_url(platform="deepseek"):
    default_url = "https://api.deepseek.com"
    if platform == "deepseek":
        return "https://api.deepseek.com"  # 默认返回DeepSeek的API地址
    elif platform == "openai":
        return "https://api.openai.com"
    else:
        print("Unsupported platform. Using default base URL.")
        return default_url  # 不支持的平台返回默认URL

# AI计算器主函数
def ai_calculator(prompt, **kwargs):
    api_key = get_api_key()  # 获取API密钥
    if not api_key:
        return "API key not found. Please set DEEPSEEK_API_KEY in the.env file."

    # 初始化OpenAI客户端
    client = OpenAI(api_key=api_key, base_url=get_base_url())
    
    # 构造计算器提示词模板
    prompts = f"""
    你是一个专业的计算器，你可以计算任何数学问题。
    请严格遵守以下规则：
    1. 只计算数学问题，不要回答非数学问题。如果用户问你非数学问题，你可以回复"您咨询的非数学问题，我拒绝回答。"。
    2. 只返回最终的计算结果，不要返回任何解释。
    3. 如果涉及小数结果，请保留小数点后四位。

    请你计算以下问题：
    {prompt}
    """

    # 调用API获取计算结果
    stream = kwargs.get("stream", True)  # 是否使用流式响应
    response = client.chat.completions.create(
        model=kwargs.get("model", "deepseek-chat"),  # 使用的模型名称
        messages=[
            {"role": "user", "content": prompts},  # 用户消息内容
        ],
        temperature=kwargs.get("temperature", 0.1),  # 温度参数控制输出随机性
        max_tokens=kwargs.get("max_tokens", 100),  # 最大token数限制
        stream=stream,  # 是否使用流式响应
    )

    print("\nAI计算器回复结果：")
    result = ""
    if stream:
        # 处理流式响应
        for chunk in response:
            chunk_content = chunk.choices[0].delta.content
            print(chunk_content, end="")
            result += chunk_content
    else:
        # 处理非流式响应
        result = response.choices[0].message.content
        print(result)

    return result  # 返回最终计算结果


if __name__ == "__main__":
    # 主程序循环
    while True:
        print("欢迎使用AI计算器！")
        print("输入 'exit' 或 'quit' 可以退出程序。")
        prompt = input("\n请输入数学问题（可以用自然语言来描述）：")
        if prompt.lower() in ["exit", "quit"]:
            break  # 退出程序
        ai_calculator(prompt)  # 调用计算器函数
        print("\n")
